#include "ws2812.h"
#include "driver/rmt_tx.h"
#include "esp_check.h"
#define LED_STRIP_RESOLUTION 100000000
static const char *TAG = "led_encoder";
struct ws2812_strip_t
{
   rmt_channel_handle_t led_chan;
   rmt_encoder_handle_t led_encoder;
   uint8_t *led_buffer;
   int led_num;
};

typedef struct 
{
    rmt_encoder_t base;
    rmt_encoder_t *bytes_encoder;
    rmt_encoder_t *copy_encoder;
    int state;
    rmt_symbol_word_t reset_code;
}rmt_led_strip_encoder_t;

static size_t rmt_encoder_led_strip(rmt_encoder_t *encoder,rmt_channel_handle_t channel,const void* primary_data,size_t data_sizze,rmt_encode_state_t *ret_state)
{
    rmt_led_strip_encoder_t *led_encoder = __containerof(encoder,rmt_led_strip_encoder_t,base);
    rmt_encoder_handle_t bytes_encoder =  led_encoder->bytes_encoder;
    rmt_encoder_handle_t copy_encoder = led_encoder->copy_encoder;
    rmt_encode_state_t session_state =RMT_ENCODING_RESET;
    rmt_encode_state_t state = RMT_ENCODING_RESET;
    size_t encoded_symbol = 0;
    switch(led_encoder->state)
    {
        case 0:
        encoded_symbol +=bytes_encoder->encode(bytes_encoder,channel,primary_data,data_sizze,&session_state);
        if(session_state & RMT_ENCODING_COMPLETE){
            led_encoder->state = 1;
        }
        if(session_state & RMT_ENCODING_MEM_FULL)
        {
            state |= RMT_ENCODING_MEM_FULL;
            goto out;
        }
        case 1:
        encoded_symbol += copy_encoder->encode(copy_encoder,channel,&led_encoder->reset_code,sizeof(&led_encoder->reset_code),&session_state);
        if(session_state & RMT_ENCODING_COMPLETE)
        {
            led_encoder->state =RMT_ENCODING_RESET;
            state |= RMT_ENCODING_COMPLETE;

        }
        if(session_state & RMT_ENCODING_MEM_FULL)
        {
            state |= RMT_ENCODING_MEM_FULL;
            goto out;
        }
    }
    out:
    *ret_state = state;
    return encoded_symbol;


}
static esp_err_t rmt_del_led_strip_encoder(rmt_encoder_t* encoder)
{
    rmt_led_strip_encoder_t *led_encoder = __containerof(encoder,rmt_led_strip_encoder_t,base);
    rmt_del_encoder(led_encoder->bytes_encoder);
    rmt_del_encoder(led_encoder->copy_encoder);
    free(led_encoder);
    return ESP_OK;

}
static esp_err_t rmt_led_strip_encoder_reset(rmt_encoder_t *encoder)
{
    rmt_led_strip_encoder_t *led_encoder = __containerof(encoder,rmt_led_strip_encoder_t,base);
    
    rmt_encoder_reset(led_encoder->bytes_encoder);
    rmt_encoder_reset(led_encoder->copy_encoder);
    led_encoder->state = RMT_ENCODING_RESET;
    return ESP_OK;

}

esp_err_t rmt_new_led_strip_encoder(rmt_encoder_handle_t *ret_encoder)
{
    esp_err_t ret = ESP_OK;

    rmt_led_strip_encoder_t *led_encoder = NULL;

    led_encoder = calloc(1,sizeof(rmt_led_strip_encoder_t));
    ESP_GOTO_ON_FALSE(led_encoder,ESP_ERR_NO_MEM,err,TAG,"no mem for strip encoder");
    led_encoder->base.encode = rmt_encoder_led_strip;
    led_encoder->base.del = rmt_del_led_strip_encoder;
    led_encoder->base.reset =rmt_led_strip_encoder_reset;

    rmt_bytes_encoder_config_t bytes_encoder_config = {
        .bit0 = {
            .level0 = 1,
            .duration0 = 0.3*LED_STRIP_RESOLUTION/1000000,
            .level1 = 0,
            .duration1 =0.9*LED_STRIP_RESOLUTION/1000000
        },
        .bit1 = {
            .level0 = 1,
            .duration0 =0.9 * LED_STRIP_RESOLUTION /1000000,
            .level1 = 0,
            .duration1 = 0.3 * LED_STRIP_RESOLUTION/1000000,
        },
        .flags.msb_first =1
    };
    rmt_new_bytes_encoder(&bytes_encoder_config,&led_encoder->bytes_encoder);

    rmt_copy_encoder_config_t rmt_copy_encoder_config_t = {};
    rmt_new_copy_encoder(&rmt_copy_encoder_config_t,&led_encoder->copy_encoder);

    uint32_t reset_ticks =LED_STRIP_RESOLUTION/1000000*50/2;
    led_encoder->reset_code = (rmt_symbol_word_t){
        .level0 = 0,
        .duration0 = reset_ticks,
        .level1 = 1,
        .duration1 = reset_ticks,
    };
    *ret_encoder = &led_encoder->base;
    return ESP_OK;

    err:
    if(led_encoder)
    {
        if(led_encoder->bytes_encoder)
        {
            rmt_del_encoder(led_encoder->bytes_encoder);
        }
        if(led_encoder->copy_encoder)
        {
            rmt_del_encoder(led_encoder->copy_encoder);
        }
        free(led_encoder);
    }
    return ret;



}
esp_err_t ws2812_init(gpio_num_t gpio,int maxled,ws2812_strip_handle_t *handle)
{
    struct ws2812_strip_t* led_handle = NULL;
    led_handle = calloc(1,sizeof(struct ws2812_strip_t));
    assert(led_handle);
    led_handle->led_buffer = calloc(1,maxled*3);
    assert(led_handle->led_buffer);
    led_handle->led_num = maxled;


    rmt_tx_channel_config_t tx_chan_config ={
        .clk_src = RMT_CLK_SRC_DEFAULT,
        .gpio_num = gpio,
        .mem_block_symbols = 64,
        .resolution_hz = LED_STRIP_RESOLUTION,
        .trans_queue_depth = 4,

    };
    ESP_ERROR_CHECK(rmt_new_tx_channel(&tx_chan_config,&led_handle));
    ESP_ERROR_CHECK(rmt_new_led_strip_encoder(&led_handle->led_encoder));
    ESP_ERROR_CHECK(rmt_enable(led_handle->led_chan));
    *handle = led_handle;

    return ESP_OK;

}

esp_err_t ws2812_write(ws2812_strip_handle_t handle,uint32_t index,uint32_t r,uint32_t g,uint32_t b)
{
    rmt_transmit_config_t tx_config={
        .loop_count=1,
        
    };
    if(index > handle->led_num)
    {
        return ESP_FAIL;
    }
    uint32_t start = index*3;
    handle->led_buffer[start+0] = g & 0xff;
    handle->led_buffer[start+1] = b & 0xff;
    handle->led_buffer[start+2] = r & 0xff;

    return rmt_transmit(handle->led_chan,handle->led_encoder,handle->led_buffer,handle->led_num*3,&tx_config);
}